﻿//
// This file is part of AVSReflector project
// Copyright (C) 2011 Yuan Pengfei
//

using System;
using System.Collections.Generic;
using System.Text;
using System.Drawing;
using System.Drawing.Imaging;
using System.Runtime.InteropServices;

namespace AVSReflector.Core
{
    public abstract class Picture : INode
    {
        // TODO: Add support for picture display extension & region extension
        public uint picture_type;
        public uint bbv_delay;
        public uint picture_distance;
        public uint progressive_frame;
        public uint picture_structure;
        public uint top_field_first;
        public uint repeat_first_field;
        public uint fixed_picture_qp;
        public uint picture_qp;
        public uint skip_mode_flag;
        public uint loop_filter_disable;
        public uint loop_filter_parameter_flag;

        // P & B only
        public uint picture_reference_flag;

        public abstract string[] GetInfo();

        public INode[] children;

        public MacroBlock[,] BlockMatrix;

        public int[,] YMatrix;
        public int[,] CbMatrix;
        public int[,] CrMatrix;

        public Bitmap ImageBlock;
        public Bitmap ImageRGB;
        public Bitmap ImageY;
        public Bitmap ImageCb;
        public Bitmap ImageCr;

        protected void GetImageBlock()
        {
            if(ImageBlock == null)
                ImageBlock = new Bitmap(BlockMatrix.GetLength(1) * 15, BlockMatrix.GetLength(0) * 15);
            Graphics g = Graphics.FromImage(ImageBlock);
            for (int i = 0; i < BlockMatrix.GetLength(0); ++i)
                for (int j = 0; j < BlockMatrix.GetLength(1); ++j)
                    DrawBlock(g, BlockMatrix[i, j], j, i);
        }

        static void DrawBlock(Graphics g, MacroBlock mb, int x, int y)
        {
            if (mb.type == MacroBlockType.I_8x8)
                FillMatrix(g, Color.Red, x * 15 + 1, y * 15 + 1, 13, 13);
            if (MacroBlock.NameOfType(mb.type).Contains("P_"))
                FillMatrix(g, Color.Blue, x * 15 + 1, y * 15 + 1, 13, 13);
            if (mb.type == MacroBlockType.P_Skip || mb.type == MacroBlockType.B_Skip)
                FillMatrix(g, Color.Gray, x * 15 + 1, y * 15 + 1, 13, 13);
            bool vertical = false, horizontal = false;
            if (MacroBlock.NameOfType(mb.type).Contains("8x8"))
                vertical = horizontal = true;
            if (MacroBlock.NameOfType(mb.type).Contains("8x16"))
                vertical = true;
            if (MacroBlock.NameOfType(mb.type).Contains("16x8"))
                horizontal = true;
            Pen pen = new Pen(Color.Black);
            g.DrawLine(pen, x * 15, y * 15, x * 15 + 14, y * 15);
            g.DrawLine(pen, x * 15, y * 15, x * 15, y * 15 + 14);
            g.DrawLine(pen, x * 15 + 14, y * 15 + 14, x * 15 + 14, y * 15);
            g.DrawLine(pen, x * 15 + 14, y * 15 + 14, x * 15, y * 15 + 14);
            if (vertical)
                g.DrawLine(pen, x * 15 + 7, y * 15, x * 15 + 7, y * 15 + 14);
            if (horizontal)
                g.DrawLine(pen, x * 15, y * 15 + 7, x * 15 + 14, y * 15 + 7);
            if (MacroBlock.NameOfType(mb.type).Contains("B_"))
            {
                if (vertical && horizontal)
                {
                    Color[] arr = new Color[4];
                    for (int i = 0; i < 4; ++i)
                    {
                        if (mb.part_type[i] == MacroBlockPartType.SB_Direct_8x8)
                            arr[i] = Color.Brown;
                        else if (mb.part_type[i] == MacroBlockPartType.SB_Bck_8x8)
                            arr[i] = Color.Green;
                        else if (mb.part_type[i] == MacroBlockPartType.SB_Fwd_8x8)
                            arr[i] = Color.Blue;
                        else if (mb.part_type[i] == MacroBlockPartType.SB_Sym_8x8)
                            arr[i] = Color.Yellow;
                    }
                    FillMatrix(g, arr[0], x * 15 + 1, y * 15 + 1, 6, 6);
                    FillMatrix(g, arr[1], x * 15 + 8, y * 15 + 1, 6, 6);
                    FillMatrix(g, arr[2], x * 15 + 1, y * 15 + 8, 6, 6);
                    FillMatrix(g, arr[3], x * 15 + 8, y * 15 + 8, 6, 6);
                }
                else if (vertical || horizontal)
                {
                    Color a = Color.White, b = Color.White;
                    string sa = MacroBlock.NameOfType(mb.type).Substring(2, 3);
                    string sb = MacroBlock.NameOfType(mb.type).Substring(6, 3);
                    if (sa == "Fwd")
                        a = Color.Blue;
                    else if (sa == "Bck")
                        a = Color.Green;
                    else if (sa == "Sym")
                        a = Color.Yellow;
                    if (sb == "Fwd")
                        b = Color.Blue;
                    else if (sb == "Bck")
                        b = Color.Green;
                    else if (sb == "Sym")
                        b = Color.Yellow;
                    if (vertical)
                    {
                        FillMatrix(g, a, x * 15 + 1, y * 15 + 1, 6, 13);
                        FillMatrix(g, b, x * 15 + 8, y * 15 + 1, 6, 13);
                    }
                    if (horizontal)
                    {
                        FillMatrix(g, a, x * 15 + 1, y * 15 + 1, 13, 6);
                        FillMatrix(g, b, x * 15 + 1, y * 15 + 8, 13, 6);
                    }
                }
                else
                {
                    if (mb.type == MacroBlockType.B_Direct_16x16)
                        FillMatrix(g, Color.Brown, x * 15 + 1, y * 15 + 1, 13, 13);
                    if (mb.type == MacroBlockType.B_Fwd_16x16)
                        FillMatrix(g, Color.Blue, x * 15 + 1, y * 15 + 1, 13, 13);
                    if (mb.type == MacroBlockType.B_Bck_16x16)
                        FillMatrix(g, Color.Green, x * 15 + 1, y * 15 + 1, 13, 13);
                    if (mb.type == MacroBlockType.B_Sym_16x16)
                        FillMatrix(g, Color.Yellow, x * 15 + 1, y * 15 + 1, 13, 13);
                }
            }
        }

        static void FillMatrix(Graphics g, Color c, int x, int y, int w, int h)
        {
            SolidBrush br = new SolidBrush(c);
            g.FillRectangle(br, x, y, w, h);
        }

        protected void GetPictures()
        {
            ImageY = new Bitmap((int)Context.currentSequence.horizontal_size, (int)Context.currentSequence.vertical_size);
            ImageCb = new Bitmap((int)Context.currentSequence.horizontal_size, (int)Context.currentSequence.vertical_size);
            ImageCr = new Bitmap((int)Context.currentSequence.horizontal_size, (int)Context.currentSequence.vertical_size);
            ImageRGB = new Bitmap((int)Context.currentSequence.horizontal_size, (int)Context.currentSequence.vertical_size);
            int half = YMatrix.GetLength(0) / 2;

            BitmapData bmpdata = ImageY.LockBits(new Rectangle(0, 0, (int)Context.currentSequence.horizontal_size, (int)Context.currentSequence.vertical_size), ImageLockMode.ReadWrite, PixelFormat.Format24bppRgb);
            if (picture_structure == 1)
            {
                for (int i = 0; i < (int)Context.currentSequence.vertical_size; ++i)
                    for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                    {
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, (byte)YMatrix[i, j]);
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, (byte)YMatrix[i, j]);
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, (byte)YMatrix[i, j]);
                    }
            }
            else
            {
                for (int i = 0; i < (int)Context.currentSequence.vertical_size; ++i)
                {
                    if (i % 2 == 0)
                    {
                        for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                        {
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, (byte)YMatrix[i / 2, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, (byte)YMatrix[i / 2, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, (byte)YMatrix[i / 2, j]);
                        }
                    }
                    else
                    {
                        for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                        {
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, (byte)YMatrix[i / 2 + half, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, (byte)YMatrix[i / 2 + half, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, (byte)YMatrix[i / 2 + half, j]);
                        }
                    }
                }
            }
            ImageY.UnlockBits(bmpdata);

            bmpdata = ImageCb.LockBits(new Rectangle(0, 0, (int)Context.currentSequence.horizontal_size, (int)Context.currentSequence.vertical_size), ImageLockMode.ReadWrite, PixelFormat.Format24bppRgb);
            if (picture_structure == 1)
            {
                for (int i = 0; i < (int)Context.currentSequence.vertical_size; ++i)
                    for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                    {
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, (byte)CbMatrix[i, j]);
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, (byte)CbMatrix[i, j]);
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, (byte)CbMatrix[i, j]);
                    }
            }
            else
            {
                for (int i = 0; i < (int)Context.currentSequence.vertical_size; ++i)
                {
                    if (i % 2 == 0)
                    {
                        for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                        {
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, (byte)CbMatrix[i / 2, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, (byte)CbMatrix[i / 2, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, (byte)CbMatrix[i / 2, j]);
                        }
                    }
                    else
                    {
                        for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                        {
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, (byte)CbMatrix[i / 2 + half, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, (byte)CbMatrix[i / 2 + half, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, (byte)CbMatrix[i / 2 + half, j]);
                        }
                    }
                }
            }
            ImageCb.UnlockBits(bmpdata);

            bmpdata = ImageCr.LockBits(new Rectangle(0, 0, (int)Context.currentSequence.horizontal_size, (int)Context.currentSequence.vertical_size), ImageLockMode.ReadWrite, PixelFormat.Format24bppRgb);
            if (picture_structure == 1)
            {
                for (int i = 0; i < (int)Context.currentSequence.vertical_size; ++i)
                    for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                    {
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, (byte)CrMatrix[i, j]);
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, (byte)CrMatrix[i, j]);
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, (byte)CrMatrix[i, j]);
                    }
            }
            else
            {
                for (int i = 0; i < (int)Context.currentSequence.vertical_size; ++i)
                {
                    if (i % 2 == 0)
                    {
                        for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                        {
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, (byte)CrMatrix[i / 2, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, (byte)CrMatrix[i / 2, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, (byte)CrMatrix[i / 2, j]);
                        }
                    }
                    else
                    {
                        for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                        {
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, (byte)CrMatrix[i / 2 + half, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, (byte)CrMatrix[i / 2 + half, j]);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, (byte)CrMatrix[i / 2 + half, j]);
                        }
                    }
                }
            }
            ImageCr.UnlockBits(bmpdata);

            bmpdata = ImageRGB.LockBits(new Rectangle(0, 0, (int)Context.currentSequence.horizontal_size, (int)Context.currentSequence.vertical_size), ImageLockMode.ReadWrite, PixelFormat.Format24bppRgb);
            if (picture_structure == 1)
            {
                for (int i = 0; i < (int)Context.currentSequence.vertical_size; ++i)
                    for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                    {
                        byte r, g, b;
                        Decoder.YCbCr2Color(YMatrix[i, j], CbMatrix[i, j], CrMatrix[i, j], out r, out g, out b);
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, b);
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, g);
                        Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, r);
                    }
            }
            else
            {
                for (int i = 0; i < (int)Context.currentSequence.vertical_size; ++i)
                {
                    if (i % 2 == 0)
                    {
                        for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                        {
                            byte r, g, b;
                            Decoder.YCbCr2Color(YMatrix[i / 2, j], CbMatrix[i / 2, j], CrMatrix[i / 2, j], out r, out g, out b);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, b);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, g);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, r);
                        }
                    }
                    else
                    {
                        for (int j = 0; j < (int)Context.currentSequence.horizontal_size; ++j)
                        {
                            byte r, g, b;
                            Decoder.YCbCr2Color(YMatrix[i / 2 + half, j], CbMatrix[i / 2 + half, j], CrMatrix[i / 2 + half, j],
                                out r, out g, out b);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3, b);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 1, g);
                            Marshal.WriteByte(bmpdata.Scan0, i * bmpdata.Stride + j * 3 + 2, r);
                        }
                    }
                }
            }
            ImageRGB.UnlockBits(bmpdata);
        }

        public void Cleanup()
        {
            YMatrix = null;
            CrMatrix = null;
            CbMatrix = null;

            for (int i = 0; i < children.Length; ++i)
            {
                if (children[i].GetType() != typeof(Slice))
                    continue;
                Slice s = (Slice)children[i];
                s.blocks = null;
            }

            for(int i=0; i<BlockMatrix.GetLength(0); ++i)
                for (int j = 0; j < BlockMatrix.GetLength(1); ++j)
                {
                    MacroBlock b = BlockMatrix[i, j];
                    b.pred_mode_flag = null;
                    b.intra_luma_pred_mode = null;
                    b.mb_reference_index = null;
                    b.mv_diff_x = null;
                    b.mv_diff_y = null;
                    b.inter_luma_pred_mode = null;
                    b.ref_index = null;
                    b.diff_x = null;
                    b.diff_y = null;
                }
        }
    }
}
